import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import networkx as nx


def convert_neighbour(connections, pos):
    # 先转换为邻接矩阵形式
    neighbour_matrix = np.ones((130, 130)) * np.inf
    for index, row in connections.iterrows():
        p1 = row[0]
        for item in row[1:]:
            if pd.isna(item):
                continue
            dist = np.sqrt((pos.loc[p1, 'x'] - pos.loc[item, 'x']) ** 2 + (pos.loc[p1, 'y'] - pos.loc[item, 'y']) ** 2)
            neighbour_matrix[pos.loc[p1, 'idx'], pos.loc[item, "idx"]] = dist
            neighbour_matrix[pos.loc[item, 'idx'], pos.loc[p1, "idx"]] = dist
    return neighbour_matrix


def draw_graph(connections, pos, routines=None, routine_color=None):
    fig, ax = plt.subplots()
    # draw edges
    edges = []
    good_points = pos.loc["J01":"J20"].index
    for index, row in connections.iterrows():
        p1 = row[0]
        for item in row[1:]:
            if pd.isna(item):
                continue
            edges.append((p1, item))
    edges = set(edges)
    for edge in edges:
        p1, p2 = edge
        color = "red" if (p1 in good_points) and (p2 in good_points) else "gray"
        ax.plot([pos.loc[p1, 'x'], pos.loc[p2, 'x']], [pos.loc[p1, "y"], pos.loc[p2, 'y']], c=color)

    # draw nodes
    d_points = pos.loc["D1": 'D2']
    z_points = pos.loc["Z01":"Z06"]
    f_points = pos.loc["F01": "F60"]
    j_points = pos.loc["J01": "J62"]
    point_sets = [d_points, z_points, f_points, j_points]
    colors = ["green", "blue", "red", "gray"]

    for point_set, color in zip(point_sets, colors):
        ax.scatter(x=point_set["x"], y=point_set["y"], c=color)

    if routines is not None:
        routine_color = "green" if routine_color is None else routine_color
        if any(isinstance(i, list) for i in routines):
            for single_routine in routines:
                for idx in range(len(single_routine) - 1):
                    p1, p2 = single_routine[idx], single_routine[idx + 1]
                    x1, y1 = pos.iloc[p1, 0], pos.iloc[p1, 1]
                    x2, y2 = pos.iloc[p2, 0], pos.iloc[p2, 1]
                    ax.plot([x1, x2], [y1, y2], routine_color)
        else:
            routine = routines
            for idx in range(len(routine)-1):
                p1, p2 = routine[idx], routine[idx+1]
                x1, y1 = pos.iloc[p1, 0], pos.iloc[p1, 1]
                x2, y2 = pos.iloc[p2, 0], pos.iloc[p2, 1]
                ax.plot([x1, x2], [y1, y2], routine_color)
    plt.show()


def shortest_path(start_point, graph, end_point=None):
    p = nx.shortest_path(graph, source=start_point, weight='weight', method='dijkstra')
    if end_point is not None:
        return p[end_point]
    return p


if __name__ == '__main__':
    pos = pd.read_excel("data.xls", sheet_name=0, index_col=0)
    connections = pd.read_excel("data.xls", sheet_name=1, header=None)
    # draw the graph
    # draw_graph(connections, pos)

    # create the matrix
    neighbour = convert_neighbour(connections, pos)
    graph = nx.from_numpy_matrix(neighbour)
    # method参数默认为'dijkstra'，当边有向时，可以选择'bellman-ford'
    # 算从D1到Z01-Z06的最短路径
    z_points = pos.loc["Z01":"Z06"]
    p = nx.shortest_path(graph, source=pos.loc["D1", 'idx'], weight='weight', method='dijkstra')
    routines = [p[pos.loc[dest, 'idx']] for dest in z_points.index]
    print('D1到Z01-Z06的最短路径为:', routines)

    draw_graph(connections, pos, routines=routines, routine_color="yellow")
    best_roads = shortest_path(0, graph)
    print(best_roads)